{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "mnist_pca_oss.ipynb",
      "provenance": [
        {
          "file_id": "1EiPzbzKpWttz_w7kidxhKA9ebTfgjFw9",
          "timestamp": 1656341553042
        },
        {
          "file_id": "1X6kerSJWNq29dFSt7wsG9xUTbaH9inPJ",
          "timestamp": 1651771834451
        },
        {
          "file_id": "1tkajifmOG3L04k-u3S9TJ3p0OTAksTPW",
          "timestamp": 1651592497504
        },
        {
          "file_id": "1JGwpGam_Sux-tUchjBQsQwZ8mO0hpfG0",
          "timestamp": 1651584778312
        },
        {
          "file_id": "1pFYW--edbeDa90WGqRYurf7qLJ-mcKOK",
          "timestamp": 1598656150914
        },
        {
          "file_id": "1DZIxYeEZARlT0ez71zbL14IpXi_r-1tb",
          "timestamp": 1589483831073
        }
      ],
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "//learning/deepmind/public/tools/ml_python:ml_notebook",
        "kind": "private"
      }
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "metadata": {
        "cellView": "form",
        "id": "xAnK07Bq0QtN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fxCOZqTSLJ26",
        "cellView": "form"
      },
      "source": [
        "# @title Imports\n",
        "from scipy.optimize import linear_sum_assignment\n",
        "from functools import partial\n",
        "import jax\n",
        "from jax import grad, jit, make_jaxpr, vmap, random, pmap\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "from scipy.spatial.distance import cosine as cosine_distance\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.cm as cm\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import time\n",
        "\n",
        "from jax.example_libraries import optimizers\n",
        "\n",
        "import tensorflow as tf\n",
        "from six.moves import cPickle as pickle #for performance\n",
        "import optax as optix\n",
        "import dill\n",
        "\n",
        "gfile = tf.io.gfile"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PJhpUcQ56-KS",
        "cellView": "form"
      },
      "source": [
        "#@title Optimizers\n",
        "\n",
        "def nesterov(eta):\n",
        "  return optix.chain(\n",
        "      optix.trace(decay=0.9, nesterov=True),\n",
        "      optix.scale(-eta))\n",
        "  \n",
        "def sgd(eta):\n",
        "  return optix.chain(\n",
        "      optix.trace(decay=0.0, nesterov=False),\n",
        "      optix.scale(-eta))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f7sT5jM4LI70",
        "cellView": "form"
      },
      "source": [
        "#@title Alternative Updates\n",
        "\n",
        "@partial(jax.jit, static_argnums=(5,6))\n",
        "def gha_update(vi, weights, V, opt_state, X, eta=1e-4, opt=sgd):\n",
        "    # vi is (d,)\n",
        "    # weights is (k,)\n",
        "    # V is (k, d), i.e., vectors on rows\n",
        "    # X is (N, d)\n",
        "    gs_ii = []\n",
        "    gs_ij = []\n",
        "    Xvi = jnp.dot(X, vi)\n",
        "    XV = jnp.transpose(jnp.dot(X, jnp.transpose(V)))  # Xvj on row j\n",
        "    for j in range(V.shape[0]):\n",
        "      viMvj = jnp.dot(Xvi, XV[j])\n",
        "      # vjMvj = jnp.dot(XV[j], XV[j])\n",
        "      gs_ii += [jnp.dot(jnp.transpose(X), jnp.dot(X, V[j]))]  # only needed for gs_ii\n",
        "      gs_ij += [viMvj * V[j]]\n",
        "    weights_ii = (jnp.sign(weights - 0.5) + 1.) / 2.  # maps 1 to 1 else to 0\n",
        "    weights_ij = (jnp.sign(weights + 0.5) - 1.) / 2.  # maps -1 to -1 else to 0\n",
        "    weights_ij -= weights_ii\n",
        "    gs_ii = jnp.dot(jnp.transpose(jnp.array(gs_ii)), weights_ii)\n",
        "    gs_ij = jnp.dot(jnp.transpose(jnp.array(gs_ij)), weights_ij)\n",
        "    # gs = gs_ii + gs_ij\n",
        "    # grads = jnp.dot(jnp.transpose(X), jnp.dot(X, g_ii)) + gs_ij\n",
        "\n",
        "    grads = gs_ii + gs_ij\n",
        "\n",
        "    # grads = jnp.dot(jnp.transpose(X), jnp.dot(X, gs))\n",
        "    # grads -= jnp.dot(grads, vi) * vi\n",
        "    \n",
        "    # This computes and applies updates with optix and updates the opt_state\n",
        "    # updates, opt_state = nesterov(eta).update(-grads, opt_state)\n",
        "    updates, opt_state = opt(eta).update(-grads, opt_state)\n",
        "    vi_new = optix.apply_updates(vi, updates)\n",
        "\n",
        "    vi_new /= jnp.clip(jnp.linalg.norm(vi_new), a_min=1.)\n",
        "    # vi_new = jnp.clip(vi_new, a_min=-1., a_max=1.)\n",
        "    return vi_new, opt_state\n",
        "\n",
        "@partial(jax.jit, static_argnums=(5,6))\n",
        "def ojas_deflation_update(vi, weights, V, opt_state, X, eta=1e-4, opt=sgd):\n",
        "    # vi is (d,)\n",
        "    # weights is (k,)\n",
        "    # V is (k, d), i.e., vectors on rows\n",
        "    # X is (N, d)\n",
        "    gs_ii = []\n",
        "    gs_ij = []\n",
        "    Xvi = jnp.dot(X, vi)\n",
        "    XV = jnp.transpose(jnp.dot(X, jnp.transpose(V)))  # Xvj on row j\n",
        "    for j in range(V.shape[0]):\n",
        "      vivj = jnp.dot(vi, V[j])\n",
        "      # vjMvj = jnp.dot(XV[j], XV[j])\n",
        "      gs_ii += [jnp.dot(jnp.transpose(X), XV[j])]  # only needed for gs_ii\n",
        "      gs_ij += [vivj * V[j]]\n",
        "    weights_ii = (jnp.sign(weights - 0.5) + 1.) / 2.  # maps 1 to 1 else to 0\n",
        "    weights_ij = (jnp.sign(weights + 0.5) - 1.) / 2.  # maps -1 to -1 else to 0\n",
        "    # weights_ij -= weights_ii\n",
        "    gs_ii = jnp.dot(jnp.transpose(jnp.array(gs_ii)), weights_ii)\n",
        "    gs_ij = jnp.dot(jnp.transpose(jnp.array(gs_ij)), weights_ij)\n",
        "    # gs = gs_ii + gs_ij\n",
        "    # grads = jnp.dot(jnp.transpose(X), jnp.dot(X, g_ii)) + gs_ij\n",
        "\n",
        "    grads = gs_ii + gs_ij\n",
        "\n",
        "    # grads = jnp.dot(jnp.transpose(X), jnp.dot(X, gs))\n",
        "    # grads -= jnp.dot(grads, vi) * vi\n",
        "    \n",
        "    # This computes and applies updates with optix and updates the opt_state\n",
        "    # updates, opt_state = nesterov(eta).update(-grads, opt_state)\n",
        "    updates, opt_state = opt(eta).update(-grads, opt_state)\n",
        "    vi_new = optix.apply_updates(vi, updates)\n",
        "\n",
        "    vi_new /= jnp.clip(jnp.linalg.norm(vi_new), a_min=1.)\n",
        "    # vi_new = jnp.clip(vi_new, a_min=-1., a_max=1.)\n",
        "    return vi_new, opt_state\n",
        "\n",
        "@partial(jit, static_argnums=(3,4))\n",
        "def matrix_krasulinas_update(V, opt_state, X, eta=1e-4, opt=sgd):\n",
        "    # vectors on rows of V\n",
        "    grad = jnp.dot(jnp.transpose(X), jnp.dot(X, jnp.transpose(V)))\n",
        "    grad -= jnp.dot(jnp.transpose(V), jnp.dot(V, grad))\n",
        "    \n",
        "    # This computes and applies updates with optix and updates the opt_state\n",
        "    updates, opt_state = opt(eta).update(-jnp.transpose(grad), opt_state)\n",
        "    Vstep = optix.apply_updates(V, updates)\n",
        "\n",
        "    Q, R = jnp.linalg.qr(jnp.transpose(Vstep))\n",
        "    signs = jnp.sign(jnp.sign(jnp.diag(R)) + .5)\n",
        "    V_new = jnp.transpose(Q * signs)\n",
        "    return V_new, opt_state\n",
        "\n",
        "@partial(jit, static_argnums=(3,4))\n",
        "def ojas_update(V, opt_state, X, eta=1e-4, opt=sgd):\n",
        "    # X is minibatch of size (batch_size, d) and d is dimensionality\n",
        "    # for mnist, d is 784.\n",
        "    # V is minibatch of size (k, d) where k is num_eigenvectors\n",
        "    # vectors on rows of V\n",
        "    grad = jnp.dot(jnp.transpose(X), jnp.dot(X, jnp.transpose(V)))\n",
        "    \n",
        "    # This computes and applies updates with optix and updates the opt_state\n",
        "    updates, opt_state = opt(eta).update(-jnp.transpose(grad), opt_state)\n",
        "    Vstep = optix.apply_updates(V, updates)\n",
        "\n",
        "    Q, R = jnp.linalg.qr(jnp.transpose(Vstep))\n",
        "    signs = jnp.sign(jnp.sign(jnp.diag(R)) + .5)\n",
        "    V_new = jnp.transpose(Q * signs)\n",
        "    return V_new, opt_state\n",
        "\n",
        "@partial(jit)\n",
        "def sherman_morrison_woodbury(Apinv, u, v):\n",
        "  # u and v should be 1-d vectors\n",
        "  uvT = jnp.outer(u, v)\n",
        "  Apinv_uvT_Apinv = jnp.dot(Apinv, jnp.dot(uvT, Apinv))\n",
        "  c = jnp.dot(v, jnp.dot(Apinv, u))\n",
        "  return Apinv - Apinv_uvT_Apinv / (1 + c)\n",
        "\n",
        "@partial(jit, static_argnums=(3,4))\n",
        "def implicit_matrix_krasulinas_update(V, Vpinv, X, eta=1e-4, opt=sgd):\n",
        "    del opt\n",
        "    # only works with mb_size = 1 !!!!\n",
        "    # vectors on rows of V (k,d)\n",
        "    yt = X\n",
        "    C = jnp.transpose(V)\n",
        "    Cpinv = jnp.transpose(Vpinv)\n",
        "\n",
        "    xt = jnp.dot(Cpinv, yt)  # X (d,), Vpinv (d,k), xt (k,)\n",
        "    rt = jnp.dot(C, xt) - yt  # rt (d,)\n",
        "    etaxt = eta / (1 + eta * jnp.sum(xt**2.))  # scalar\n",
        "    deltat = -etaxt * jnp.outer(rt, xt)  # deltat (d,k)\n",
        "\n",
        "    Vdeltat = jnp.transpose(deltat)\n",
        "\n",
        "    V_new = V + Vdeltat\n",
        "\n",
        "    # Cpinv_new = (C - etaxt * outer(rt, xt))^{-1}\n",
        "    Cpinv_new = sherman_morrison_woodbury(Cpinv, -etaxt * rt, xt)\n",
        "    Vpinv_new = jnp.transpose(Cpinv_new)\n",
        "\n",
        "    return V_new, Vpinv_new"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lObamxqKJR8e",
        "cellView": "form"
      },
      "source": [
        "#@title Data Class Util\n",
        "\n",
        "class Data(object):\n",
        "  def __init__(self, ds, mb_size, k, num_samples=None,\n",
        "               one_device=False, shuffle_data=True, center=False, unit_var=False,\n",
        "               svd_by_evd=False, data_stats_dict=None):\n",
        "\n",
        "    self.mb_size = mb_size\n",
        "    self.ds = ds\n",
        "\n",
        "    self.k = k\n",
        "    self.shuffle_data = shuffle_data\n",
        "    self.one_device = one_device\n",
        "\n",
        "    self.num_devices = jax.local_device_count()\n",
        "    self.batch_dims = self.get_batch_dims(k, self.one_device)\n",
        "\n",
        "    self.X = None\n",
        "    self.generator = self.make_generator(self.ds, self.batch_dims,\n",
        "                                         self.shuffle_data)\n",
        "\n",
        "    if svd_by_evd:\n",
        "      self.mean = 0\n",
        "      Cov = 0\n",
        "      batch_size = np.prod(self.batch_dims)\n",
        "      num_batches, remainder = divmod(num_samples, batch_size)\n",
        "      mean = 0\n",
        "      if remainder > 0:\n",
        "        print(\"Warning: mb_size does not evenly divide dataset\")\n",
        "        print(\"num_samples, num_batches, remainder\",\n",
        "              num_samples, num_batches, remainder)\n",
        "      if center:\n",
        "        for i in range(num_batches):\n",
        "          if i % (num_batches // 10) == 0:\n",
        "            print(\"{:d} / {:d}\".format(i, num_batches), flush=True)\n",
        "          xi = jnp.reshape(self.reshape_data(next(self.generator)),\n",
        "                          (batch_size, -1))\n",
        "          mean += xi / float(num_samples)\n",
        "      for i in range(num_batches):\n",
        "        if i % (num_batches // 10) == 0:\n",
        "          print(\"{:d} / {:d}\".format(i, num_batches), flush=True)\n",
        "        xi = jnp.reshape(self.reshape_data(next(self.generator)),\n",
        "                         (batch_size, -1))\n",
        "        Cov += jnp.dot(jnp.transpose(xi - mean), xi - mean)\n",
        "      self.mean = mean\n",
        "      dims = Cov.shape[0]\n",
        "      print(\"Computing principal components...\", flush=True)\n",
        "      Sigma2, V = np.linalg.eigh(np.array(Cov))  # returns in ascending order\n",
        "      # change to descending order\n",
        "      Sigma2 = Sigma2[::-1]\n",
        "      V = V[:, ::-1]\n",
        "      Sigma = np.sqrt(np.real(Sigma2))\n",
        "      Vh = V.T\n",
        "    else:\n",
        "      X = self.load_from_ds(ds, num_samples).astype(np.float32)\n",
        "      self.mean = np.zeros(X.shape[1])\n",
        "      if center:\n",
        "        self.mean = np.mean(X, axis=0)\n",
        "        X -= self.mean\n",
        "      self.std = np.ones(X.shape[1])\n",
        "      if unit_var:\n",
        "        self.std = np.std(X, axis=0)\n",
        "        X /= self.std\n",
        "\n",
        "      num_samples = X.shape[0]\n",
        "      if len(X.shape[1:]) > 1:\n",
        "        X = np.reshape(X, (num_samples, -1))\n",
        "      dims = X.shape[1]\n",
        "      print(\"Computing principal components...\")\n",
        "      _, Sigma, Vh = np.linalg.svd(X, full_matrices=False,\n",
        "                                  compute_uv=True)\n",
        "      Cov = np.dot(X.T, X) / X.shape[0]\n",
        "      self.X = X\n",
        "\n",
        "    self.num_samples = num_samples\n",
        "    self.dims = dims\n",
        "    self.Sigma = Sigma\n",
        "    self.Vh = Vh\n",
        "    self.Cov = Cov\n",
        "\n",
        "  def load_from_ds(self, ds, num_samples):\n",
        "    if isinstance(ds, np.ndarray):\n",
        "      return np.reshape(ds, (ds.shape[0], -1))\n",
        "    else:\n",
        "      if num_samples is None:\n",
        "        err_msg = (\"num_samples must be passed to Data constructor \"\n",
        "                   \"if using tfds.\")\n",
        "        raise ValueError(err_msg)\n",
        "      X = next(self.make_generator(ds, [num_samples], shuffle_data=False))\n",
        "      X = np.reshape(X['image'], (num_samples, -1)) / 255.\n",
        "      return X\n",
        "\n",
        "  def get_batch_dims(self, k, one_device=False):\n",
        "    if one_device:\n",
        "      self.k_per_device = k\n",
        "      return [self.mb_size]\n",
        "    else:\n",
        "      self.k_per_device = k // self.num_devices\n",
        "      err_msg = f\"specify a k that num_devices={self.num_devices} divides evenly\"\n",
        "      if self.num_devices * self.k_per_device != k:\n",
        "        raise ValueError(err_msg)\n",
        "      return [self.num_devices, self.k_per_device, self.mb_size]\n",
        "\n",
        "  def set_multi_device(self, multi):\n",
        "    if self.one_device == multi:\n",
        "      self.one_device = not self.one_device\n",
        "      self.batch_dims = self.get_batch_dims(self.k, self.one_device)\n",
        "      self.generator = self.make_generator(self.ds, self.batch_dims,\n",
        "                                           self.shuffle_data)\n",
        "\n",
        "  def set_k(self, k):\n",
        "    self.k = k\n",
        "    self.batch_dims = self.get_batch_dims(k, self.one_device)\n",
        "    self.generator = self.make_generator(self.ds, self.batch_dims,\n",
        "                                         self.shuffle_data)\n",
        "  \n",
        "  def set_mb_size(self, mb_size):\n",
        "    self.mb_size = mb_size\n",
        "    self.batch_dims = self.get_batch_dims(self.k, self.one_device)\n",
        "    self.generator = self.make_generator(self.ds, self.batch_dims,\n",
        "                                         self.shuffle_data)\n",
        "\n",
        "  def make_generator(self, ds, batch_dims, shuffle_data, seed=None):\n",
        "    \"\"\"Loads the dataset as a generator of batches.\"\"\"\n",
        "    if ds is None:\n",
        "      inds = np.arange(self.num_samples)\n",
        "      if shuffle_data:\n",
        "        np.random.seed(seed)\n",
        "        np.random.shuffle(inds)\n",
        "      num_batches = self.num_samples // self.mb_size\n",
        "      num_samples_even = num_batches * self.mb_size\n",
        "      X = np.reshape(self.X[inds[:num_samples_even]], [num_batches] + batch_dims + [-1])\n",
        "      yield from np.split(X, num_batches, axis=0)\n",
        "    else:\n",
        "      total_batch_size = jnp.prod(jnp.array(batch_dims))\n",
        "      if shuffle_data:\n",
        "          ds = ds.shuffle(10 * total_batch_size, seed=seed)\n",
        "      for batch_size in reversed(batch_dims):\n",
        "        ds = ds.batch(batch_size)\n",
        "      yield from tfds.as_numpy(ds)\n",
        "\n",
        "  def reset_generator(self, seed=12345):\n",
        "    self.generator = self.make_generator(self.ds, self.batch_dims, self.shuffle_data, seed)\n",
        "  \n",
        "  def reshape_data(self, data_input):\n",
        "    # expected data_input shape = (num_devices, k_per_device, mb_size, *data_shape)\n",
        "    if 'image' in data_input:\n",
        "      data_input = jnp.array(data_input['image']) / 255.\n",
        "    data_input_dims = len(data_input.shape)\n",
        "    data_input = data_input.astype(np.float32)\n",
        "    expected_dims = self.one_device * 2 + (1 - self.one_device) * 4\n",
        "    if len(data_input.shape) > expected_dims:\n",
        "      dim_size = np.prod(data_input.shape[expected_dims - 1:])\n",
        "      shape = list(data_input.shape[:expected_dims - 1]) + [dim_size]\n",
        "      return (jnp.reshape(data_input, shape) - self.mean) / self.std\n",
        "    else:\n",
        "      return (jnp.array(data_input) - self.mean) / self.std"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l_YgzbzZrqcI",
        "cellView": "form"
      },
      "source": [
        "#@title Metrics Class Util\n",
        "\n",
        "def eval_to_class(ylim=None, bucket=None):\n",
        "  def inner(evaluator):\n",
        "    class Eval(object):\n",
        "      def __init__(self, evaluator, ylim, bucket):\n",
        "        self.evaluator = evaluator\n",
        "        self.__str__ = evaluator.__name__\n",
        "        self.ylim = ylim\n",
        "        self.bucket = bucket\n",
        "      \n",
        "      def __call__(self, V, **kwargs):\n",
        "        return self.evaluator(V, **kwargs)\n",
        "    return Eval(evaluator, ylim, bucket)\n",
        "  return inner\n",
        "\n",
        "class Metrics(object):\n",
        "  def __init__(self, evaluators, **kwargs):\n",
        "    eval_names = [e.__str__ for e in evaluators]\n",
        "    self.evaluators = dict(zip(eval_names, evaluators))\n",
        "    self.match_eigvecs = kwargs[\"match_eigvecs\"]  # not used\n",
        "    self.fixed_kwargs = kwargs\n",
        "    self.V_last = None\n",
        "\n",
        "    if (\"X\" in kwargs and kwargs[\"X\"] is not None) and (\"Cov\" in kwargs):\n",
        "      X = kwargs[\"X\"]\n",
        "      Ustar = kwargs[\"Ustar\"]\n",
        "      top1 = kwargs[\"Vh\"][0]\n",
        "      Ustar1 = jnp.outer(top1, top1)\n",
        "      min_recon_error = jnp.linalg.norm(X - jnp.dot(X, Ustar))\n",
        "      top1_recon_error = jnp.linalg.norm(X - jnp.dot(X, Ustar1))\n",
        "      min_compression_loss = jnp.trace(jnp.dot(jnp.eye(Ustar.shape[0]) - Ustar, kwargs[\"Cov\"]))\n",
        "      top1_compression_loss = jnp.trace(jnp.dot(jnp.eye(Ustar1.shape[0]) - Ustar1, kwargs[\"Cov\"]))\n",
        "      ev_error_scale = jnp.linalg.norm(Ustar1 - Ustar)\n",
        "\n",
        "      self.fixed_kwargs.update({\"min_recon_error\": min_recon_error,\n",
        "                                \"top1_recon_error\": top1_recon_error,\n",
        "                                \"min_compression_loss\": min_compression_loss,\n",
        "                                \"top1_compression_loss\": top1_compression_loss,\n",
        "                                \"ev_error_scale\": ev_error_scale})\n",
        "    \n",
        "    empty_lists = [[] for _ in evaluators]\n",
        "    self.progress = []\n",
        "    self.record = dict(zip(eval_names, empty_lists))\n",
        "  \n",
        "  def evaluate(self, iteration, epoch, time, V):\n",
        "    # V contains vectors on rows\n",
        "    # Q, R = jnp.linalg.qr(jnp.transpose(V))\n",
        "    # signs = jnp.sign(jnp.sign(jnp.diag(R)) + .5)\n",
        "    # Vorth = jnp.transpose(Q * signs)\n",
        "    # P = jnp.dot(Vorth.transpose(), Vorth)\n",
        "    self.progress.append((iteration, epoch, time))\n",
        "    Vpinv = jnp.array(np.linalg.pinv(np.array(V).T))\n",
        "    P = jnp.dot(V.transpose(), Vpinv)\n",
        "    for key, evaluator in self.evaluators.items():\n",
        "      value = evaluator(V, P=P, V_last=self.V_last,\n",
        "                        **self.fixed_kwargs)\n",
        "      self.record[key].append(value)\n",
        "    if 'norm_of_drift' in self.evaluators:\n",
        "      self.V_last = V\n",
        "  \n",
        "  def update(self, key, value):\n",
        "    if key in self.record:\n",
        "      self.record[key].append(value)\n",
        "    else:\n",
        "      self.record[key] = [value]\n",
        "\n",
        "def sine_squared(u, v):\n",
        "  cosine = 1 - cosine_distance(u, v)\n",
        "  return 1. - cosine**2.\n",
        "\n",
        "@eval_to_class()\n",
        "def compression_loss(V, P, Cov, min_compression_loss,\n",
        "                     top1_compression_loss, **kwargs):\n",
        "  del kwargs\n",
        "  compress_loss = jnp.trace(jnp.dot(jnp.eye(P.shape[0]) - P, Cov))\n",
        "  scale = top1_compression_loss - min_compression_loss\n",
        "  return (compress_loss - min_compression_loss) / scale\n",
        "\n",
        "@eval_to_class()\n",
        "def recon_error(V, P, X, min_recon_error, top1_recon_error,\n",
        "                **kwargs):\n",
        "  del kwargs\n",
        "  est_recon = jnp.dot(X, P)\n",
        "  recon_err = jnp.linalg.norm(X - est_recon)\n",
        "  scale = top1_recon_error - min_recon_error\n",
        "  return (recon_err - min_recon_error) / scale\n",
        "\n",
        "@eval_to_class()\n",
        "def ev_error(V, Ustar, P, ev_error_scale, **kwargs):\n",
        "  del kwargs\n",
        "  k = V.shape[0]\n",
        "  eigvec_err = jnp.linalg.norm(P - Ustar)\n",
        "  return eigvec_err / ev_error_scale\n",
        "\n",
        "@eval_to_class(ylim=[0,1])\n",
        "def ev_error_to_v1(V, Vh, **kwargs):\n",
        "  del kwargs\n",
        "  eigvec_err = jnp.array([sine_squared(vi, Vh[0]) for vi in V])\n",
        "  return eigvec_err\n",
        "\n",
        "@eval_to_class(ylim=[0,90], bucket=4)\n",
        "def individual_ev_error_bucket(V, Vh, match_eigvecs=False, **kwargs):\n",
        "  del kwargs\n",
        "  if match_eigvecs:\n",
        "    errs = np.zeros((len(V), len(Vh)))\n",
        "    for i, vi in enumerate(V):\n",
        "      for j, vh in enumerate(Vh):\n",
        "        errs[i, j] = sine_squared(vi, vh)\n",
        "    row_ind, col_ind = linear_sum_assignment(errs)\n",
        "    eigvec_err = errs[row_ind, col_ind]\n",
        "  else:\n",
        "    eigvec_err = jnp.array([sine_squared(vi, vh) for vi, vh in zip(V, Vh)])\n",
        "  return jnp.arcsin(jnp.sqrt(eigvec_err)) * 180. / np.pi\n",
        "\n",
        "@eval_to_class(ylim=[0,90], bucket=None)\n",
        "def individual_ev_error_deg(V, Vh, match_eigvecs=False, **kwargs):\n",
        "  del kwargs\n",
        "  if match_eigvecs:\n",
        "    errs = np.zeros((len(V), len(Vh)))\n",
        "    for i, vi in enumerate(V):\n",
        "      for j, vh in enumerate(Vh):\n",
        "        errs[i, j] = sine_squared(vi, vh)\n",
        "    row_ind, col_ind = linear_sum_assignment(errs)\n",
        "    eigvec_err = errs[row_ind, col_ind]\n",
        "  else:\n",
        "    eigvec_err = jnp.array([sine_squared(vi, vh) for vi, vh in zip(V, Vh)])\n",
        "  return jnp.arcsin(jnp.sqrt(eigvec_err)) * 180. / np.pi\n",
        "\n",
        "@eval_to_class()\n",
        "def longest_streak_of_correct_eigvecs(V, Vh, **kwargs):\n",
        "  del kwargs\n",
        "  eigvec_err = jnp.array([sine_squared(vi, vh) for vi, vh in zip(V, Vh)])\n",
        "  radians = jnp.arcsin(jnp.sqrt(eigvec_err))\n",
        "  learned = (radians <= np.pi / 8)\n",
        "  num_learned = np.sum(np.cumsum(learned) >= np.arange(1, len(learned) + 1))\n",
        "  return num_learned\n",
        "\n",
        "@eval_to_class()\n",
        "def neurips_loss(V, Ustar, P, **kwargs):\n",
        "  del kwargs\n",
        "  k = V.shape[0]\n",
        "  neurips_loss = 1 - np.trace(Ustar.dot(P)) / float(k)\n",
        "  return neurips_loss\n",
        "\n",
        "@eval_to_class()\n",
        "def ortho_error(V, **kwargs):\n",
        "  del kwargs\n",
        "  k = V.shape[0]\n",
        "  sqrt_num_pairs = np.sqrt(k * float(k - 1.))\n",
        "  orth_err = jnp.linalg.norm(jnp.dot(V, V.transpose()) - jnp.eye(V.shape[0]))\n",
        "  return orth_err / sqrt_num_pairs\n",
        "\n",
        "@eval_to_class()\n",
        "def norm_of_drift(V, V_last, **kwargs):\n",
        "  del kwargs\n",
        "  if V_last is None:\n",
        "    return np.NaN * jnp.ones(V.shape[0])\n",
        "  else:\n",
        "    return jnp.linalg.norm(V - V_last, axis=1)\n",
        "\n",
        "evaluators = [recon_error, ev_error, longest_streak_of_correct_eigvecs, individual_ev_error_bucket, individual_ev_error_deg, neurips_loss, compression_loss, ortho_error, ev_error_to_v1]\n",
        "evaluators_short = [longest_streak_of_correct_eigvecs, individual_ev_error_bucket, individual_ev_error_deg, neurips_loss, norm_of_drift]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NmRbNanG41lW",
        "cellView": "form"
      },
      "source": [
        "#@title QR Experiment Util\n",
        "\n",
        "def qr_experiment(V0, k, lr, num_epochs, update,\n",
        "                  optimizer, evaluators, data, match_eigvecs=False,\n",
        "                  zero_mean=False, unit_variance=False):  \n",
        "  Ustar = np.dot(data.Vh[:k].T, data.Vh[:k])\n",
        "  metrics = Metrics(evaluators, X=data.X, Vh=data.Vh, Cov=data.Cov,\n",
        "                    Ustar=Ustar, match_eigvecs=match_eigvecs)\n",
        "\n",
        "  key = jax.random.PRNGKey(1234)\n",
        "  V = jax.random.normal(key, (k, data.dims))\n",
        "  V = V/jnp.linalg.norm(V, axis=1, keepdims=True)\n",
        "  \n",
        "  weights = get_weights(k).reshape(k, -1)\n",
        "  opt_state = optimizer(lr).init(V)\n",
        "\n",
        "  epochs_to_iters = (data.num_samples / data.mb_size)\n",
        "  num_iters = int(num_epochs * epochs_to_iters)\n",
        "\n",
        "  start = time.time()\n",
        "\n",
        "  for i in range(num_iters + 1):\n",
        "    # retrieve minibatch\n",
        "    minibatch = data.reshape_data(next(data.generator))\n",
        "\n",
        "    if i % (num_iters // 10) == 0:\n",
        "      # evaluate current vectors\n",
        "      epoch = i / epochs_to_iters\n",
        "      print(\"# of epochs = {:.1f}, # of iterations = {:d}\".format(epoch, i))\n",
        "\n",
        "      metrics.evaluate(i, epoch, time.time() - start, V)\n",
        "\n",
        "    # update vectors\n",
        "    V, opt_state = update(V.reshape(k, -1), opt_state, minibatch, lr,\n",
        "                          optimizer)\n",
        "\n",
        "  return V.reshape((k, -1)), metrics"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Lissa Experiment Util\n",
        "\n",
        "def lissa_experiment(V0, k, lr, num_epochs, update,\n",
        "                  optimizer, evaluators, data, match_eigvecs=False,\n",
        "                  zero_mean=False, unit_variance=False):  \n",
        "  Ustar = np.dot(data.Vh[:k].T, data.Vh[:k])\n",
        "  metrics = Metrics(evaluators, X=data.X, Vh=data.Vh, Cov=data.Cov,\n",
        "                    Ustar=Ustar, match_eigvecs=match_eigvecs)\n",
        "\n",
        "  key = jax.random.PRNGKey(1234)\n",
        "  V = jax.random.normal(key, (k, data.dims))\n",
        "  V = V/jnp.linalg.norm(V, axis=1, keepdims=True)\n",
        "  \n",
        "  weights = get_weights(k).reshape(k, -1)\n",
        "  opt_state = optimizer(lr).init(V.T)\n",
        "\n",
        "  epochs_to_iters = (data.num_samples / data.mb_size)\n",
        "  num_iters = int(num_epochs * epochs_to_iters)\n",
        "\n",
        "  start = time.time()\n",
        "  key = jax.random.PRNGKey(42)\n",
        "\n",
        "  for i in range(num_iters + 1):\n",
        "    # retrieve minibatch\n",
        "    # 32 x 784\n",
        "    minibatch = data.reshape_data(next(data.generator))\n",
        "    # 8 x 784\n",
        "    # 1 x 784 --> fair comparisons\n",
        "\n",
        "    # batch_size x 128 -- sample rows from minibatch\n",
        "    # sampled_minibatch = np.random.choice()\n",
        "    # 1 x 128\n",
        "    # 16 x 784\n",
        "\n",
        "    if i % (num_iters // 10) == 0:\n",
        "      # evaluate current vectors\n",
        "      epoch = i / epochs_to_iters\n",
        "      print(\"# of epochs = {:.1f}, # of iterations = {:d}\".format(epoch, i))\n",
        "\n",
        "      metrics.evaluate(i, epoch, time.time() - start, V)\n",
        "\n",
        "    # update vectors\n",
        "    V, opt_state, key = update(V.reshape(k, -1), opt_state, minibatch, key, lr,\n",
        "                          optimizer)\n",
        "\n",
        "  return V.reshape((k, -1)), metrics\n",
        "\n",
        "\n",
        "# @partial(jit, static_argnums=(4, 5, 6, 7))\n",
        "def lissa_update(V, opt_state, X, key, eta=1e-4, opt=sgd,\n",
        "                estimator='lissa', p=0.1):\n",
        "    # (32 x 784)\n",
        "    # (16 X 784)\n",
        "    # X is minibatch of size (batch_size, d) and d is dimensionality\n",
        "    # for mnist, d is 784.\n",
        "    # V is minibatch of size (k, d) where k is num_eigenvectors\n",
        "    # vectors on rows of V\n",
        "\n",
        "    # N X D \n",
        "    # N states, D tasks\n",
        "    # Output: N x d\n",
        "\n",
        "    # 60000 X 784\n",
        "    # 16 x 784\n",
        "    # N = 784, D = 60000\n",
        "    # ----- ----- ----- -----\n",
        "    # Phi: N = 784, batch_size = 32\n",
        "    # Psi: N = 784, d = 16\n",
        "\n",
        "    # 784 x 60000 \n",
        "    # 784 x 32\n",
        "    # 2 x 1\n",
        "    # num_rows = 3\n",
        "    # 784 X 60000 <--\n",
        "\n",
        "    # S x T\n",
        "    # Original matrix: 784 x 60000\n",
        "    # Sample rows: 784 x 2\n",
        "    ##### Other methods #####\n",
        "    # 784 x 1\n",
        "    ## Our method can do 784 x 1 or smaller (128 x 1)\n",
        "\n",
        "  \n",
        "\n",
        "    key, subkey, subkey2 = jax.random.split(key, 3)\n",
        "    # iteration1 = jax.random.randint(subkey2, shape=(), minval=1, maxval=10)\n",
        "    iteration1 = 10\n",
        "    print(iteration1)\n",
        "    # iteration1 = jnp.ceil(jnp.log(jax.random.uniform(subkey2, shape=())) / jnp.log1p(-p))\n",
        "    # iteration1 = iteration1.astype('int')\n",
        "    k, d = V.shape\n",
        "    Phi, opt_state, _ = nabla_phi_analytical( \n",
        "        Phi=V.T,  # (d x k) d=784, \n",
        "        Psi=X.T,  # (d x batch_size) d =784,\n",
        "        key=subkey,\n",
        "        optim=opt(eta),\n",
        "        opt_state=opt_state,\n",
        "        estimator=estimator,\n",
        "        alpha=0.9,\n",
        "        use_l2_reg=False,\n",
        "        reg_coeff=False,\n",
        "        use_penalty=False,\n",
        "        # j=d,\n",
        "        j=iteration1,\n",
        "        num_rows=d, # This is to be changed.\n",
        "    )\n",
        "    V_new = Phi.T\n",
        "    return V_new, opt_state, key"
      ],
      "metadata": {
        "id": "qECd_qM1S56D",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Lissa code\n",
        "\n",
        "import optax\n",
        "import functools\n",
        "\n",
        "def matrix_estimator(Phi, num_rows, key):  # pylint: disable=invalid-name\n",
        "  r\"\"\"Computes an unbiased estimate of an input matrix.\n",
        "\n",
        "  $\\nu(s_i)^{-1}e_{s_i} \\phi_{s_i}^\\T$\n",
        "\n",
        "  Args:\n",
        "    Phi: S times d array\n",
        "    num_rows: int: number of rows used in estimators\n",
        "    key: prng key\n",
        "  Returns:\n",
        "    S times d array\n",
        "  \"\"\"\n",
        "  S, _ = Phi.shape  # pylint: disable=invalid-name\n",
        "  states = jax.random.randint(key, (num_rows,), 0, S)\n",
        "  # states = jax.random.permutation(key, jnp.arange(S))[:num_rows]\n",
        "  mask = jnp.zeros_like(Phi)\n",
        "  mask = mask.at[states].set(1)\n",
        "  return Phi * mask\n",
        "\n",
        "@jax.jit\n",
        "def _russian_roulette(Phi, states, coefficients, alpha):  # pylint: disable=invalid-name\n",
        "  \"\"\"Computes Russian roulette given fixed number of iterations.\"\"\"\n",
        "  S, d = Phi.shape  # pylint: disable=invalid-name\n",
        "  I = jnp.eye(d)  # pylint: disable=invalid-name\n",
        "\n",
        "  def _lissa_body(carry, state):\n",
        "    lissa_j = alpha * I\n",
        "    lissa_j += (\n",
        "        I - alpha * S * jnp.einsum('i,j->ij', Phi[state], Phi[state])) @ carry\n",
        "    return lissa_j, lissa_j\n",
        "\n",
        "  lissa_init = alpha * I\n",
        "  _, lissa_estimates = jax.lax.scan(_lissa_body, lissa_init, states)\n",
        "\n",
        "  deltas = I - S * jnp.einsum('ni,nj,mjk->mik', Phi[states], Phi[states],\n",
        "                              lissa_estimates)\n",
        "  deltas *= coefficients.reshape(-1, 1, 1)\n",
        "\n",
        "  return jnp.sum(deltas, axis=0)\n",
        "\n",
        "\n",
        "def russian_roulette(Phi, p, key, coeff_alpha):  # pylint: disable=invalid-name\n",
        "  \"\"\"Computes the Russian roulette estimator from a LISSA sequence.\n",
        "\n",
        "  Args:\n",
        "    Phi: S times d array\n",
        "    p: paramter of the bernoulli distribution\n",
        "    key: prng key\n",
        "    coeff_alpha: float\n",
        "  Returns:\n",
        "    array of shape d times d\n",
        "  \"\"\"\n",
        "  S, _ = Phi.shape  # pylint: disable=invalid-name\n",
        "  norm = jnp.linalg.norm(Phi.T @ Phi, ord=2)\n",
        "  alpha = coeff_alpha * 1 / norm\n",
        "\n",
        "  # Sample from geometric R.V. to get number of iterations\n",
        "  key, subkey = jax.random.split(key)\n",
        "  iterations = int(\n",
        "      jnp.ceil(jnp.log(jax.random.uniform(subkey)) / jnp.log1p(-p)))\n",
        "\n",
        "  # Sample states\n",
        "  key, subkey = jax.random.split(key)\n",
        "  states = jax.random.randint(subkey, (iterations,), 0, S)\n",
        "\n",
        "  # Get delta coefficients\n",
        "  coefficients = alpha / ((1 - p) ** jnp.arange(1, iterations + 1))\n",
        "\n",
        "  return _russian_roulette(Phi, states, coefficients, alpha)\n",
        "\n",
        "@functools.partial(jax.jit, static_argnums=(1, 2, 4, 5, 6, 7))\n",
        "def lissa(Phi,  # pylint: disable=invalid-name\n",
        "          j,\n",
        "          num_rows,\n",
        "          key,\n",
        "          coeff_alpha,\n",
        "          use_penalty=False,\n",
        "          reg_coeff=0.0):\n",
        "  \"\"\"Computes the lissa estimator.\n",
        "\n",
        "  Args:\n",
        "    Phi: S times d array\n",
        "    j: int, index of the lissa estimator\n",
        "    num_rows: int: number of rows used in estimators\n",
        "    key: prng key\n",
        "    coeff_alpha: float\n",
        "    use_penalty: bool: whether to add \"lambda * Id\" term to features\n",
        "    reg_coeff: float: coeff for reg\n",
        "    normalization: bool: whether to use S to normalize alpha\n",
        "\n",
        "  Returns:\n",
        "    d times d array\n",
        "  \"\"\"\n",
        "  num_rows = 1\n",
        "  S, d = Phi.shape  # pylint: disable=invalid-name\n",
        "  I = jnp.eye(d)  # pylint: disable=invalid-name\n",
        "\n",
        "  def _neumann_series(carry, state):\n",
        "    A_j = alpha * I  # pylint: disable=invalid-name\n",
        "    # pylint: disable=invalid-name\n",
        "    if use_penalty:\n",
        "      A_j += (I - alpha * (1 / num_rows) *\n",
        "              (Phi[state, :].T @ Phi[state, :] + reg_coeff * I)) @ carry\n",
        "    else:\n",
        "      A_j += (I - alpha *\n",
        "              (1 / num_rows) * Phi[state, :].T @ Phi[state, :]) @ carry\n",
        "    return A_j, None\n",
        "\n",
        "  _, subkey = jax.random.split(key)\n",
        "  states = jax.random.randint(subkey, (j, 1), 0, S)\n",
        "  # pdb.set_trace()\n",
        "  norm = 2 * jnp.max(jnp.sum(jnp.square(Phi[states.reshape(j, )]), axis=1))\n",
        "  # norm = jnp.linalg.norm(Phi[states.reshape(j, )].T @ Phi[states.reshape(j, )]  , ord=2) / num_rows\n",
        "  # norm = jnp.linalg.norm(Phi.T @ Phi , ord=2) / num_rows\n",
        "  alpha = coeff_alpha * 1 / norm\n",
        "  # alpha = 2.\n",
        "  lissa_init = alpha * I\n",
        "  lisa_j, _ = jax.lax.scan(_neumann_series, lissa_init, states)\n",
        "  return lisa_j\n",
        "\n",
        "def least_square_estimator(Phi,  # pylint: disable=invalid-name\n",
        "                           Psi,\n",
        "                           num_rows,\n",
        "                           j,\n",
        "                           key,\n",
        "                           estimator='lissa',\n",
        "                           alpha=0.9,\n",
        "                           use_penalty=False,\n",
        "                           reg_coeff=0.0):  # pylint: disable=invalid-name\n",
        "  r\"\"\"Computes an unbiased least squares estimate.\n",
        "\n",
        "  $W^*_\\Phi = (\\Phi^T \\Phi)^{-1} \\Phi^T \\Psi$\n",
        "\n",
        "  Args:\n",
        "    Phi: S times d array\n",
        "    Psi: S times T array\n",
        "    num_rows: int: number of rows used in estimators\n",
        "    j: int: num of samples for lissa\n",
        "    key: prng key\n",
        "    estimator: str: russian_roulette, lissa, hat_w\n",
        "    alpha: float: renormalize covariance term\n",
        "    use_penalty: bool: whether to add \"lambda * Id\" term to features\n",
        "    reg_coeff: float: coeff for reg\n",
        "\n",
        "  Returns:\n",
        "    array d times T\n",
        "  \"\"\"\n",
        "  S, _ = Phi.shape  # pylint: disable=invalid-name\n",
        "  key, subkey = jax.random.split(key)\n",
        "  states = jax.random.randint(subkey, (num_rows,), 0, S)\n",
        "  _, subkey = jax.random.split(key)\n",
        "  if estimator == 'lissa':\n",
        "    cov_estim = lissa(Phi, j, num_rows, subkey, alpha, use_penalty, reg_coeff)\n",
        "  # cov_estim = jnp.linalg.solve(Phi.T @ Phi, jnp.eye(d))\n",
        "  return cov_estim @ Phi[states, :].T @ Psi[\n",
        "      states, :] / num_rows  # we use the same samples here\n",
        "\n",
        "@functools.partial(jax.jit, static_argnums=(3, 5, 6, 7, 8, 9, 10, 11, 12))\n",
        "def nabla_phi_analytical(  # pylint: disable=invalid-name\n",
        "    Phi,\n",
        "    Psi,\n",
        "    key,\n",
        "    optim,\n",
        "    opt_state,\n",
        "    estimator,\n",
        "    alpha,\n",
        "    use_l2_reg,\n",
        "    reg_coeff,\n",
        "    use_penalty,\n",
        "    j,\n",
        "    num_rows=1):\n",
        "\n",
        "  r\"\"\"Computes unbiased estimate of 2 * (\\Phi W^*_\\Phi - \\Psi)(W^*_\\Phi)^T.\n",
        "\n",
        "  Args:\n",
        "    Phi: S times d array\n",
        "    Psi: S times T array\n",
        "    key: prng key\n",
        "    optim: optax optimizer\n",
        "    opt_state: optimizer initialization\n",
        "    estimator: str: russian_roulette, lissa, hat_w\n",
        "    alpha: float: used to nornalize covariance term\n",
        "    use_l2_reg: bool: whether to use l2 reg\n",
        "    reg_coeff: float: coeff for reg\n",
        "    use_penalty: bool: whether to add \"lambda * Id\" term to features\n",
        "    j: int: num of samples for lissa\n",
        "    num_rows: int: number of rows used in estimators\n",
        "\n",
        "  Returns:\n",
        "    array S times d\n",
        "  \"\"\"\n",
        "  key, subkey = jax.random.split(key)\n",
        "  S, T = Psi.shape  # pylint: disable=invalid-name\n",
        "  task = jax.random.randint(subkey, (1,), 0, T)\n",
        "  key, subkey = jax.random.split(key)\n",
        "  Phi_estim = matrix_estimator(Phi, num_rows, subkey)  # pylint: disable=invalid-name\n",
        "  Psi_estim = matrix_estimator(Psi[:, task], num_rows, subkey)  # pylint: disable=invalid-name\n",
        "  key, subkey = jax.random.split(key)\n",
        "  least_square_estim_1 = least_square_estimator(Phi, Psi[:, task], num_rows, j,\n",
        "                                                subkey, estimator, alpha,\n",
        "                                                use_penalty, reg_coeff)\n",
        "\n",
        "  key, subkey = jax.random.split(key)\n",
        "  least_square_estim_2 = least_square_estimator(Phi, Psi[:, task], num_rows, j,\n",
        "                                                subkey, estimator, alpha,\n",
        "                                                use_penalty, reg_coeff)\n",
        "  grads = (Phi_estim @ least_square_estim_1 -\n",
        "           Psi_estim) @ least_square_estim_2.T\n",
        "  if use_l2_reg:\n",
        "    grads += Phi_estim * reg_coeff\n",
        "  # if jnp.linalg.norm(grads) > 100000.0:\n",
        "  #   import pdb\n",
        "  #   pdb.set_trace()\n",
        "  updates, opt_state = optim.update(grads, opt_state, Phi)\n",
        "  # beta = 1 / (1 + 0.1 * epoch)\n",
        "  beta = 1\n",
        "  \n",
        "  return optax.apply_updates(Phi, beta * updates), opt_state, grads\n",
        "\n",
        "def train(Phi, Psi, num_epochs, learning_rate, key, estimator, alpha,  # pylint: disable=invalid-name\n",
        "          optimizer, use_l2_reg, reg_coeff, use_penalty, j, num_rows):\n",
        "  \"\"\"Training function.\"\"\"\n",
        "  Phis = [Phi]  # pylint: disable=invalid-name\n",
        "  grads = []\n",
        "  if optimizer == 'sgd':\n",
        "    optim = optax.sgd(learning_rate)\n",
        "  elif optimizer == 'adam':\n",
        "    optim = optax.adam(learning_rate)\n",
        "  opt_state = optim.init(Phi)\n",
        "  for _ in tqdm(range(num_epochs)):\n",
        "    key, subkey = jax.random.split(key)\n",
        "    Phi, opt_state, grad = nabla_phi_analytical(\n",
        "        Phi, Psi, subkey, optim, opt_state, estimator, alpha, use_l2_reg,\n",
        "        reg_coeff, use_penalty, j, num_rows)\n",
        "    Phis.append(Phi)\n",
        "    grads.append(grad)\n",
        "  return jnp.stack(Phis), jnp.stack(grads)\n"
      ],
      "metadata": {
        "id": "_iVuo4a6VXXt",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# TODO \n",
        "\n",
        "Charline\n",
        "1. Run with batch size of 1 (number of columns = 1) and compare with other methods -- 1 x 784 matrices \n",
        "2. Run with batch size of 1 and smaller number of rows (d = 128 or smaller) and show similar performance to 1."
      ],
      "metadata": {
        "id": "WDgQGtwmorrr"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EMhZtZ5jwe4W",
        "cellView": "form"
      },
      "source": [
        "#@title General Experiment Util\n",
        "\n",
        "def init_V(num_devices, k_per_device, dims):\n",
        "  keys = jax.random.split(jax.random.PRNGKey(1234), num_devices)\n",
        "  V = jax.pmap(lambda key: jax.random.normal(key, (k_per_device, dims)))(keys)\n",
        "  V = jax.pmap(lambda V: V/jnp.linalg.norm(V, axis=1, keepdims=True))(V)\n",
        "  return V\n",
        "\n",
        "def get_weights(k):\n",
        "  num_devices = jax.local_device_count()\n",
        "  k_per_device = k // num_devices\n",
        "  weights = np.eye(k) * 2 - np.ones((k, k))\n",
        "  weights[np.triu_indices(k, 1)] = 0.\n",
        "  weights_jnp = jnp.array(weights)\n",
        "  weights_jnp = jnp.reshape(weights_jnp, [num_devices, k_per_device, k])\n",
        "  return weights_jnp\n",
        "\n",
        "def run_sweep(experiment_runner, V0, k, lr, num_epochs, update,\n",
        "              optimizer, evaluators, data, num_trials=10, match_eigvecs=False):\n",
        "    sweep_record = dict()\n",
        "    border = \"*\" * 20\n",
        "    avg_time = 0.0\n",
        "    data.reset_generator()\n",
        "    for t in range(num_trials):\n",
        "      print(\"\\n{:s}TRIAL = {:d}{:s}\\n\".format(border, t, border))\n",
        "      start = time.time()\n",
        "      _, metrics = experiment_runner(\n",
        "        V0, k, lr, num_epochs, update=update, optimizer=optimizer,\n",
        "        evaluators=evaluators, data=data, match_eigvecs=match_eigvecs)\n",
        "      end = time.time()\n",
        "      avg_time += (end - start)\n",
        "      for key, value in metrics.record.items():\n",
        "        if key in sweep_record:\n",
        "          sweep_record[key] += [value]\n",
        "        else:\n",
        "          sweep_record[key] = [value]\n",
        "    avg_time /= float(num_trials)\n",
        "    return metrics.progress, sweep_record, avg_time\n",
        "\n",
        "def run_lr_sweep(data, V0, k, num_epochs, lrs, num_trials=10):\n",
        "  data.set_multi_device(False)\n",
        "  # names = [\"Ojas\", \"EigenGame\", \"GHA\", \"EGGHA\", \"MatKras\"]\n",
        "  names = [\"Ojas\"]\n",
        "  empties = [dict() for _ in names]\n",
        "  xs_lrs = dict(zip(names, empties))\n",
        "\n",
        "  for loglr in lrs:\n",
        "    lr = 10.**float(loglr)\n",
        "    print(\"#\" * 50)\n",
        "    print(\"LEARNING RATE = {:g}\".format(lr))\n",
        "    print(\"#\" * 50)\n",
        "\n",
        "    print(\"[START - 0] ojas\" + \"\\t\"*5)\n",
        "    ojas_x, ojas_sweep, ojas_time = run_sweep(qr_experiment,\n",
        "    V0, k, lr, num_epochs, update=ojas_update, optimizer=sgd,\n",
        "    evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    xs_lrs[\"Ojas\"].update({str(loglr): [ojas_x, ojas_sweep, ojas_time]})\n",
        "    print(\"[END - 0] ojas\" + \"\\t\"*5, ojas_time)\n",
        "\n",
        "    # print(\"[START - 1] eigengame\" + \"\\t\"*5)\n",
        "    # eigengame_x, eigengame_sweep, eigengame_time = run_sweep(game_experiment,\n",
        "    #     V0, k, lr, num_epochs, update=update, optimizer=sgd,\n",
        "    #     evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    # xs_lrs[\"EigenGame\"].update({str(loglr): [eigengame_x, eigengame_sweep, eigengame_time]})\n",
        "    # print(\"[END - 1] eigengame\" + \"\\t\"*5, eigengame_time)\n",
        "\n",
        "    # print(\"[START - 2] gha\" + \"\\t\"*5)\n",
        "    # gha_x, gha_sweep, gha_time = run_sweep(game_experiment,\n",
        "    # V0, k, lr, num_epochs, update=gha_update, optimizer=sgd,\n",
        "    # evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    # xs_lrs[\"GHA\"].update({str(loglr): [gha_x, gha_sweep, gha_time]})\n",
        "    # print(\"[END - 2] gha\" + \"\\t\"*5, gha_time)\n",
        "\n",
        "    # print(\"[START - 3] eggha\" + \"\\t\"*5)\n",
        "    # eggha_x, eggha_sweep, eggha_time = run_sweep(game_experiment,\n",
        "    # V0, k, lr, num_epochs, update=eggha_update, optimizer=sgd,\n",
        "    # evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    # xs_lrs[\"EGGHA\"].update({str(loglr): [eggha_x, eggha_sweep, eggha_time]})\n",
        "    # print(\"[END - 3] eggha\" + \"\\t\"*5, eggha_time)\n",
        "\n",
        "    # print(\"[START - 4] matkras\" + \"\\t\"*5)\n",
        "    # matkras_x, matkras_sweep, matkras_time = run_sweep(qr_experiment,\n",
        "    # V0, k, lr, num_epochs, update=matrix_krasulinas_update, optimizer=sgd,\n",
        "    # evaluators=evaluators_short, data=data, num_trials=num_trials, match_eigvecs=True)\n",
        "    # xs_lrs[\"MatKras\"].update({str(loglr): [matkras_x, matkras_sweep, matkras_time]})\n",
        "    # print(\"[END - 4] matkras\" + \"\\t\"*5, matkras_time)\n",
        "  \n",
        "  return xs_lrs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Lissa Sweep\n",
        "\n",
        "def run_lissa_sweep(data, V0, k, num_epochs, lrs, num_trials=10):\n",
        "  data.set_multi_device(False)\n",
        "  # names = [\"Ojas\", \"EigenGame\", \"GHA\", \"EGGHA\", \"MatKras\"]\n",
        "  names = [\"Lissa\"]\n",
        "  empties = [dict() for _ in names]\n",
        "  xs_lrs = dict(zip(names, empties))\n",
        "\n",
        "  for loglr in lrs:\n",
        "    lr = 10.**float(loglr)\n",
        "    print(\"#\" * 50)\n",
        "    print(\"LEARNING RATE = {:g}\".format(lr))\n",
        "    print(\"#\" * 50)\n",
        "\n",
        "    print(\"[START - 0] Lissa\" + \"\\t\"*5)\n",
        "    ojas_x, ojas_sweep, ojas_time = run_sweep(lissa_experiment,\n",
        "    V0, k, lr, num_epochs, update=lissa_update, optimizer=optax.adam,\n",
        "    evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    xs_lrs[\"Lissa\"].update({str(loglr): [ojas_x, ojas_sweep, ojas_time]})\n",
        "    print(\"[END - 0] Lissa\" + \"\\t\"*5, ojas_time)\n",
        "  \n",
        "  return xs_lrs\n",
        "\n",
        "\n",
        "def run_rr_sweep(data, V0, k, num_epochs, lrs, num_trials=10):\n",
        "  data.set_multi_device(False)\n",
        "  # names = [\"Ojas\", \"EigenGame\", \"GHA\", \"EGGHA\", \"MatKras\"]\n",
        "  names = [\"RR\"]\n",
        "  empties = [dict() for _ in names]\n",
        "  xs_lrs = dict(zip(names, empties))\n",
        "\n",
        "  update = functools.partial(lissa_update, estimator='russian_roulette')\n",
        "\n",
        "  for loglr in lrs:\n",
        "    lr = 10.**float(loglr)\n",
        "    print(\"#\" * 50)\n",
        "    print(\"LEARNING RATE = {:g}\".format(lr))\n",
        "    print(\"#\" * 50)\n",
        "\n",
        "    print(\"[START - 0] Lissa\" + \"\\t\"*5)\n",
        "    ojas_x, ojas_sweep, ojas_time = run_sweep(lissa_experiment,\n",
        "    V0, k, lr, num_epochs, update=update, optimizer=optax.adam,\n",
        "    evaluators=evaluators_short, data=data, num_trials=num_trials)\n",
        "    xs_lrs[\"RR\"].update({str(loglr): [ojas_x, ojas_sweep, ojas_time]})\n",
        "    print(\"[END - 0] Lissa\" + \"\\t\"*5, ojas_time)\n",
        "  \n",
        "  return xs_lrs"
      ],
      "metadata": {
        "id": "xvxrx2DJcdNA",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "- *run_lr_sweep* returns a dictionary mapping algorithm names (hard coded within run_lr_sweep at top) to results dictionaries\n",
        "\n",
        "- each algorithm's results dictionary contains learning rates as keys mapping to values [metrics.progress --> see Metrics class, sweep_record, avg runtime (scalar) over trials]\n",
        "\n",
        "- sweep_record: dictionary mapping metric names to lists of values over training, e.g., xs_lrs['EGGHAgrad']['-4'][1]['individual_ev_error_deg'] is of shape num_trials x (num_epochs + 1) x k"
      ],
      "metadata": {
        "id": "z_YW0JcQNDWx"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cGuVnxcR4itQ"
      },
      "source": [
        "ds = tfds.load(\"mnist:3.*.*\", split=\"train\").cache().repeat()\n",
        "num_samples = 60000"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QlGT6ziq9Iow"
      },
      "source": [
        "# MNIST\n",
        "k = 16\n",
        "V0 = None\n",
        "num_epochs = 10\n",
        "mb_size = 256\n",
        "mb_size_for_svd = 15\n",
        "svd_by_evd = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mr2af6PIxxU7"
      },
      "source": [
        "data = Data(ds, mb_size_for_svd, k, num_samples=num_samples, svd_by_evd=svd_by_evd, center=True, unit_var=False, one_device=True)\n",
        "print(\"Computed principal components.\")\n",
        "data.set_k(k)\n",
        "data.set_mb_size(mb_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# S X T\n",
        "(784, 60000)\n",
        "\n",
        "# S X D\n",
        "(784, 16)"
      ],
      "metadata": {
        "id": "nVNQd8UrTnHt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "data.set_mb_size(60000)"
      ],
      "metadata": {
        "id": "QYPAAg6eac-K"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "num_epochs = 10000"
      ],
      "metadata": {
        "id": "-QjU866skvWH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "lrs = [-1]  # specified in log_10 space\n",
        "xs_lrs_lissa = run_lissa_sweep(data, V0, k, num_epochs, lrs, num_trials=1)"
      ],
      "metadata": {
        "id": "a-fUfVXYcw7E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "lrs = [-1]  # specified in log_10 space\n",
        "xs_lrs_lissa = run_rr_sweep(data, V0, k, num_epochs, lrs, num_trials=1)"
      ],
      "metadata": {
        "id": "McXTcn_PGN-Y"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GFEH20fJ-QXL"
      },
      "source": [
        "lrs = [-3]  # specified in log_10 space\n",
        "xs_lrs_oja = run_lr_sweep(data, V0, k, num_epochs, lrs, num_trials=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "xs_lrs_oja['Ojas']['-3'][1]"
      ],
      "metadata": {
        "id": "tRsQKjFaZXwl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "to_plot = xs_lrs_lissa['Lissa']['-1'][1]['neurips_loss'][0]\n",
        "plt.plot(to_plot, color='blue')\n",
        "to_plot = xs_lrs_oja['Ojas']['-3'][1]['neurips_loss'][0]\n",
        "plt.plot(to_plot, color='red')"
      ],
      "metadata": {
        "id": "yU6r5OW1mGUm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Batch size\n",
        "\n",
        "xs_lrs['Ojas']['-3'][1]"
      ],
      "metadata": {
        "id": "dTajxs5ja6ed"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6y59lzbU0I7s",
        "cellView": "form"
      },
      "source": [
        "#@title Save and Collect Results\n",
        "\n",
        "base_file_path = \"/tmp/\"\n",
        "\n",
        "def save_data(data, name):\n",
        "  def _disabled_seek(*_):\n",
        "    raise AttributeError('seek() is disabled on this object.')\n",
        "  filename = base_file_path + name + '.pkl'\n",
        "  with gfile.GFile(filename, 'wb') as out_f:\n",
        "    setattr(out_f, 'seek', _disabled_seek)\n",
        "    pickle.dump(data, out_f)\n",
        "\n",
        "def load_data(file_path):\n",
        "  with gfile.GFile(file_path, 'rb') as fin:\n",
        "    return dill.load(fin)\n",
        "\n",
        "def collect_best_results(xs_lrs, metric, methods, sign=1.):\n",
        "  print(methods)\n",
        "  method_map = dict(zip(methods, np.arange(len(methods))))\n",
        "  best_lrs = [np.nan for _ in methods]\n",
        "  best_xs = [np.nan for _ in methods]\n",
        "  best_sweeps = [np.nan for _ in methods]\n",
        "  for method, xs_lr in xs_lrs.items():\n",
        "    best_lr = -np.inf\n",
        "    best_score = -np.inf\n",
        "    for lr, data in xs_lr.items():\n",
        "      score = np.mean(data[1][metric], axis=0)[-1]\n",
        "      print(\"current best {} vs candidate {}\".format(best_score, sign * score))\n",
        "      if sign * score > best_score:\n",
        "        best_score = sign * score\n",
        "        best_lr = lr\n",
        "    print(\"method_map[method] = {}\".format(method_map[method]))\n",
        "    print(\"best lr = {}\".format(best_lr))\n",
        "    best_lrs[method_map[method]] = best_lr\n",
        "    best_xs[method_map[method]] = xs_lr[best_lr][0]\n",
        "    best_sweeps[method_map[method]] = xs_lr[best_lr][1]\n",
        "  return best_lrs, best_xs, best_sweeps\n",
        "\n",
        "def collect_lr_results(xs_lrs, metric, methods, lrs):\n",
        "  method_map = dict(zip(methods, np.arange(len(methods))))\n",
        "  xs = [np.nan for _ in methods]\n",
        "  sweeps = [np.nan for _ in methods]\n",
        "  for method, xs_lr in xs_lrs.items():\n",
        "    lr = lrs[method_map[method]]\n",
        "    xs[method_map[method]] = xs_lr[lr][0]\n",
        "    sweeps[method_map[method]] = xs_lr[lr][1]\n",
        "  return lrs, xs, sweeps"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FH7jBvCkBNx5"
      },
      "source": [
        "save_data(xs_lrs, \"mnist\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7hvKBy0VqkQf",
        "cellView": "form"
      },
      "source": [
        "#@title Plotting Util\n",
        "\n",
        "axis_label_fs = 18\n",
        "tick_label_fs = 18\n",
        "title_fs = 18\n",
        "legend_fs = 18\n",
        "lw = 4\n",
        "\n",
        "def plot_statistic(ax, x, data, legend_label, color,\n",
        "                   lw_mean, lw_fill, fill_alpha=0.5,\n",
        "                   log_scale_y=False, legend_format=\"{:.1f}\",\n",
        "                   ls='-', skip=1):\n",
        "  x = x[::skip]\n",
        "  mean = np.mean(data, axis=0)[::skip]\n",
        "  sem = np.std(data, axis=0)[::skip] / np.sqrt(len(data))\n",
        "  # sem = np.std(data, axis=0)\n",
        "  if legend_format is None:\n",
        "    legend_label_with_val = legend_label\n",
        "  else:\n",
        "    legend_label_with_val = legend_label + \": \" + legend_format.format(mean[-1])\n",
        "  if log_scale_y:\n",
        "    ax.semilogy(x, mean, color=color, ls=ls, lw=lw_mean, label=legend_label_with_val)\n",
        "  else:\n",
        "    ax.plot(x, mean, color=color, ls=ls, lw=lw_mean, label=legend_label_with_val)\n",
        "  ax.fill_between(x, mean - sem, mean + sem, lw=lw_fill, color=color,\n",
        "                  alpha=fill_alpha)\n",
        "\n",
        "def set_plot_frame(ax, xlabel, ylabel, title, ylim, incl_legend=True,\n",
        "                   label_locs=None, legend_labels=None,\n",
        "                   legend_colors=None, excluded_methods=None,\n",
        "                   num_epochs=50, mb_size=1024, num_samples=60000,\n",
        "                   x2factor=1e6, x2label='millions', decimals=0,\n",
        "                   plot_x_axis=True):\n",
        "  ax.set_ylabel(ylabel, fontsize=axis_label_fs)\n",
        "  ax.axes.tick_params(labelsize=tick_label_fs)\n",
        "  ax.set_ylim(ylim)\n",
        "  ax.set_title(title, fontsize=title_fs)\n",
        "  if label_locs:\n",
        "    for i, (label_loc, legend_label, color) in enumerate(zip(label_locs, legend_labels, legend_colors)):\n",
        "      if i in excluded_methods:\n",
        "        continue\n",
        "      ax.text(*label_loc, legend_label, fontsize=int(legend_fs*.75),\n",
        "              color=color,\n",
        "              bbox=dict(boxstyle='round', facecolor='white',\n",
        "                        edgecolor=color, alpha=0.75))\n",
        "  elif incl_legend:\n",
        "    # ax.legend(prop={\"size\": legend_fs})\n",
        "    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5),\n",
        "              prop={\"size\": legend_fs})\n",
        "  if not plot_x_axis:\n",
        "    return\n",
        "  ax.set_xlabel(xlabel, fontsize=axis_label_fs)\n",
        "  ax2 = ax.twiny()\n",
        "  plt.gcf().subplots_adjust(bottom=0.5)\n",
        "  # Move twinned axis ticks and label from top to bottom\n",
        "  ax2.xaxis.set_ticks_position(\"bottom\")\n",
        "  ax2.xaxis.set_label_position(\"bottom\")\n",
        "  # Offset the twin axis below the host\n",
        "  ax2.spines[\"bottom\"].set_position((\"axes\", -0.4))\n",
        "  # Turn on the frame for the twin axis, but then hide all \n",
        "  # but the bottom spine\n",
        "  ax2.set_frame_on(True)\n",
        "  ax2.patch.set_visible(False)\n",
        "  # as @ali14 pointed out, for python3, use this\n",
        "  # for sp in ax2.spines.values():\n",
        "  # and for python2, use this\n",
        "  for sp in ax2.spines.values():\n",
        "    sp.set_visible(False)\n",
        "  ax2.spines[\"bottom\"].set_visible(True)\n",
        "  x = ax.get_xticks()\n",
        "  ax2.set_xticks(x)\n",
        "  conversion = num_epochs * num_samples / float(mb_size)\n",
        "  if decimals == 0:\n",
        "    ax2.set_xticklabels([int(xi * conversion / x2factor) for xi in x])\n",
        "  else:\n",
        "    ax2.set_xticklabels([np.round(xi * conversion / x2factor, decimals=decimals) for xi in x])\n",
        "  ax2.set_xlim(ax.get_xlim())\n",
        "  ax2.axes.tick_params(labelsize=tick_label_fs)\n",
        "  ax2.set_xlabel(\"Iterations ({:s})\".format(x2label), fontsize=axis_label_fs)\n",
        "\n",
        "def plot_metrics(metrics, title, mean_instead=False,\n",
        "                 separate_plots=False, inc_ylabel=True):\n",
        "  num_metrics = len(metrics.record)\n",
        "  if separate_plots:\n",
        "    axs = [plt.subplots()[1] for _ in metrics.record.keys()]\n",
        "  else:\n",
        "    fig, axs = plt.subplots(1, num_metrics, figsize=(5 * num_metrics, 4))\n",
        "    plt.suptitle(title, fontsize=title_fs)\n",
        "  for i, (metric, data) in enumerate(metrics.record.items()):\n",
        "    x = metrics.epochs\n",
        "    ylabel = metric.replace(\"_\", \" \").title()\n",
        "    if \"Longest\" in ylabel:\n",
        "      ylabel = \"Longest Correct\\nEigenvector Streak\"\n",
        "      axs[i].set_ylim([0, 16])\n",
        "    elif ylabel == \"Neurips Loss\":\n",
        "      ylabel = \"Normalized\\nRank Coverage\"\n",
        "    if inc_ylabel:\n",
        "      axs[i].set_ylabel(ylabel, fontsize=axis_label_fs)\n",
        "    if separate_plots:\n",
        "      axs[i].set_title(title, fontsize=title_fs)\n",
        "    if metric in metrics.evaluators and metrics.evaluators[metric].ylim:\n",
        "        axs[i].set_ylim(metrics.evaluators[metric].ylim)\n",
        "    if len(data[0].shape) < 1:\n",
        "      if mean_instead:\n",
        "        axs[i].plot(x, np.mean(data) * np.ones_like(x), \"-o\",\n",
        "                    label=\"Final: {:g}\".format(data[-1]), lw=lw,\n",
        "                    color=\"blue\")\n",
        "      else:\n",
        "        axs[i].plot(x, data, \"-o\",\n",
        "                    label=\"Final: {:g}\".format(data[-1]), lw=lw,\n",
        "                    color=\"blue\")\n",
        "      if metric in metrics.evaluators and metrics.evaluators[metric].ylim:\n",
        "        axs[i].set_ylim(metrics.evaluators[metric].ylim)\n",
        "      axs[i].legend(prop={\"size\": legend_fs})\n",
        "    else:\n",
        "      stacked = jnp.transpose(jnp.stack(data, axis=0))  # num vecs x num iters\n",
        "      label = \"V\"\n",
        "      if metric in metrics.evaluators:\n",
        "        ylim = metrics.evaluators[metric].ylim\n",
        "        bucket = metrics.evaluators[metric].bucket\n",
        "        if ylim and bucket:\n",
        "          if inc_ylabel:\n",
        "            axs[i].set_ylabel(\"Percentiles\")\n",
        "          ptp = ylim[1] - ylim[0]\n",
        "          width = float(ptp) / float(bucket)\n",
        "          intervals = ylim[0] + jnp.cumsum(jnp.ones(bucket) * width)  # num buckets\n",
        "          intervals = jnp.reshape(intervals, (1, 1, bucket))\n",
        "          stacked = jnp.reshape(stacked, stacked.shape + (1,))\n",
        "          bucketted = (stacked < intervals) * (stacked >= (intervals - width))\n",
        "          stacked = jnp.transpose(jnp.mean(bucketted, axis=0))\n",
        "          label = \"%%\"\n",
        "          axs[i].set_ylim([0, 1])\n",
        "      for j, datum in enumerate(stacked):\n",
        "        color_idx = float(j) / stacked.shape[0]\n",
        "        color = plt.get_cmap(\"gray\")(color_idx)\n",
        "        if mean_instead:\n",
        "          axs[i].plot(x, np.mean(datum) * np.ones_like(x), \"-o\", color=color,\n",
        "                      label=r\"${:s}_{{{:d}}}$\".format(label, j), lw=lw)\n",
        "        else:\n",
        "          axs[i].plot(x, datum, \"-o\", color=color,\n",
        "                      label=r\"${:s}_{{{:d}}}$\".format(label, j), lw=lw)\n",
        "      # axs[i].legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
        "      # axs[i].legend(loc='upper center', bbox_to_anchor=(0.5, -0.05), ncol=8)\n",
        "    axs[i].set_xlabel(\"Epochs\", fontsize=axis_label_fs)\n",
        "    axs[i].axes.tick_params(labelsize=tick_label_fs)\n",
        "    axs[i].grid()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UDBfzLeuz6LO"
      },
      "source": [
        "mb = mb_size\n",
        "xs_lrs_loaded = load_data('<Specify file path here>')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_zjVhvLwKanW"
      },
      "source": [
        "xlabel = \"Epochs\"\n",
        "title = \"MNIST (Minibatch = {})\".format(mb)\n",
        "lw_mean = 4\n",
        "lw_fill = 1\n",
        "ls = '-'\n",
        "incl_legend = True\n",
        "\n",
        "x2factor = 1e3\n",
        "x2label = 'thousands'\n",
        "decimals = 0\n",
        "\n",
        "methods = [\"EigenGame\", \"GHA\", \"MatKras\", \"Ojas\", \"EGGHA\"]\n",
        "legend_loc_mapping = [0, 1, 2, 3, 4]\n",
        "legend_labels_base = [r\"$\\alpha$-EG\", \"GHA\", \"Krasulinas\", \"Ojas\", r\"$\\mu$-EG\"]\n",
        "colors = [\"blue\", \"purple\", \"green\", \"black\", \"red\"]\n",
        "excluded_methods = []\n",
        "\n",
        "included_methods = [i for i in range(len(methods)) if i not in excluded_methods]\n",
        "colors = [colors[i] for i in included_methods]\n",
        "\n",
        "use_lrs = ['-4', '-4', '-4', '-4', '-4']\n",
        "\n",
        "key = \"longest_streak_of_correct_eigvecs\"\n",
        "best_lrs, xs, sweeps = collect_best_results(xs_lrs_loaded, key, methods, sign=1.)\n",
        "print(\"best_lrs streak\", best_lrs)\n",
        "best_lrs, xs, sweeps = collect_lr_results(xs_lrs_loaded, key, methods, use_lrs)\n",
        "legend_labels = []\n",
        "for idx, (label_base, method, lr) in enumerate(zip(legend_labels_base, methods, use_lrs)):\n",
        "  if idx not in excluded_methods:\n",
        "    if not isinstance(xs_lrs_loaded[method][lr][-1], float):\n",
        "      legend_labels += [label_base]\n",
        "      continue\n",
        "    legend_labels += [label_base + \" ({:.0f})\".format(xs_lrs_loaded[method][lr][-1])]\n",
        "ylabel = \"Longest Correct\\nEigenvector Streak\"\n",
        "\n",
        "ylim = [0, 17]\n",
        "fig, axs = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(6, 6))\n",
        "axs[0].set_yticks([0, 8, 16])\n",
        "axs[0].set_yticklabels([0, 8, 16])\n",
        "\n",
        "\n",
        "label_locs = [(-1,1), (-1,1), (-1,1), (-1,1), (-1,1)]\n",
        "if legend_loc_mapping:\n",
        "  label_locs = [label_locs[llm] for llm in legend_loc_mapping]\n",
        "\n",
        "xs = [xs[i] for i in included_methods]\n",
        "sweeps = [sweeps[i] for i in included_methods]\n",
        "for i, (color, legend_label, x, sweep) in enumerate(zip(colors, legend_labels, xs, sweeps)):\n",
        "  plot_statistic(axs[0], np.array(x)[:,1], sweep[key], legend_label,\n",
        "                 color, lw_mean, lw_fill, fill_alpha=0.5,\n",
        "                 legend_format=None, ls=ls)\n",
        "\n",
        "set_plot_frame(axs[0], xlabel, ylabel, title, ylim, incl_legend=True,\n",
        "               label_locs=[label_locs[i] for i in included_methods],\n",
        "               legend_labels=legend_labels,\n",
        "               legend_colors=colors, excluded_methods=[],\n",
        "               num_epochs=num_epochs, mb_size=mb, num_samples=num_samples,\n",
        "               x2factor=x2factor, x2label=x2label, decimals=decimals,\n",
        "               plot_x_axis=False)\n",
        "fig.tight_layout(rect=[0, 0.02, 1, 1])\n",
        "path = \"ls_k16_mb.pdf\"\n",
        "fig.savefig(path, dpi=300)\n",
        "%download_file ls_k16_mb.pdf\n",
        "\n",
        "key = \"neurips_loss\"\n",
        "best_lrs, xs, sweeps = collect_best_results(xs_lrs_loaded, key, methods, sign=-1.)\n",
        "print(\"best_lrs subspace_dist\", best_lrs)\n",
        "best_lrs, xs, sweeps = collect_lr_results(xs_lrs_loaded, key, methods, use_lrs)\n",
        "ylabel = \"Subspace Distance\"\n",
        "ylim = [0, 1]\n",
        "\n",
        "\n",
        "label_locs = [(-1,10**(-1)), (-1,10**(-1)), (-1,10**(-1)), (-1,10**(-1)), (-1,10**(-1))]\n",
        "if legend_loc_mapping:\n",
        "  label_locs = [label_locs[llm] for llm in legend_loc_mapping]\n",
        "\n",
        "xs = [xs[i] for i in included_methods]\n",
        "sweeps = [sweeps[i] for i in included_methods]\n",
        "\n",
        "for i, (color, legend_label, x, sweep) in enumerate(zip(colors, legend_labels, xs, sweeps)):\n",
        "  if color == \"gray\":\n",
        "    lsi = '--'\n",
        "  else:\n",
        "    lsi = ls\n",
        "  plot_statistic(axs[1], np.array(x)[:,1], sweep[key], legend_label,\n",
        "                 color, lw_mean, lw_fill, fill_alpha=0.5,\n",
        "                 log_scale_y=True, legend_format=None, ls=lsi)\n",
        "\n",
        "set_plot_frame(axs[1], xlabel, ylabel, \"\", ylim, incl_legend=incl_legend,\n",
        "               label_locs=[label_locs[i] for i in included_methods],\n",
        "               legend_labels=legend_labels,\n",
        "               legend_colors=colors, excluded_methods=[],\n",
        "               num_epochs=num_epochs, mb_size=mb, num_samples=num_samples,\n",
        "               x2factor=x2factor, x2label=x2label, decimals=decimals)\n",
        "\n",
        "fig.tight_layout(rect=[0.04, 0.04, 1, 1])\n",
        "path = \"nl_k16_mb.pdf\"\n",
        "fig.savefig(path, dpi=300)\n",
        "%download_file nl_k16_mb.pdf"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}